package drds.server.execute_engine;

import java.nio.ByteBuffer;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import drds.server.api_warpper.Lock;
import drds.server.net.packet.EofPacket;
import drds.server.net.packet.ResultSetHeaderPacket;
import drds.server.net.packet.RowDataPacket;
import drds.server.server.Session;
import drds.server.server.SessionContext;

public class TaskContext {
	public static final Logger LOGGER = LoggerFactory.getLogger(TaskContext.class);
	private final Tasks tasks;
	private AtomicInteger taskId = new AtomicInteger(0);
	AtomicInteger packetId = new AtomicInteger(0);
	private final Session session;
	private AtomicBoolean finished = new AtomicBoolean(false);
	private AllTaskFinishedListener allTaskFinishedListener;
	private AtomicBoolean headerWrited = new AtomicBoolean();
	private final Lock writeLock = new Lock();

	public TaskContext(Session session) {
		this.tasks = new Tasks();
		this.session = session;
	}

	public byte incrementPackageId() {
		return (byte) packetId.incrementAndGet();
	}

	public void executeSequnceTask(String[] dataNodes, String sql, TaskHandler taskHandler) {
		for (String dataNode : dataNodes) {
			Task task = new Task(taskId.incrementAndGet(), sql, dataNode, taskHandler, this);
			tasks.addTask(task, false);

		}
	}

	public Lock getWriteLock() {
		return writeLock;
	}

	public void setAllTaskFinishedListener(AllTaskFinishedListener allJobFinishedListener) {
		this.allTaskFinishedListener = allJobFinishedListener;
	}

	/**
	 * Parall并行
	 */
	public void executeParallTask(String[] dataNodes, String sql, TaskHandler taskHandler) {
		for (String dataNode : dataNodes) {
			Task task = new Task(taskId.incrementAndGet(), sql, dataNode, taskHandler, this);
			tasks.addTask(task, true);

		}
	}

	/**
	 * set no more jobs created
	 */
	public void endTaskInput() {
		tasks.setNoMoreTaskInput();
	}

	public void writeHeader(List<byte[]> leftFieldBytesList, List<byte[]> rightFieldBytesList) {
		if (headerWrited.compareAndSet(false, true)) {
			try {
				writeLock.lock();
				// write new header
				ResultSetHeaderPacket resultSetHeaderPacket = new ResultSetHeaderPacket();
				resultSetHeaderPacket.fieldCount = leftFieldBytesList.size() + rightFieldBytesList.size() - 1;
				resultSetHeaderPacket.packetId = incrementPackageId();
				LOGGER.debug("packge id " + resultSetHeaderPacket.packetId);
				SessionContext sessionContext = session.getSessionContext();
				ByteBuffer byteBuffer = resultSetHeaderPacket.write(sessionContext.allocate(), sessionContext, true);
				// wirte a fields
				for (byte[] leftFieldBytes : leftFieldBytesList) {
					leftFieldBytes[3] = incrementPackageId();
					byteBuffer = sessionContext.writeToByteBuffer(leftFieldBytes, byteBuffer);
				}
				// write b field
				for (int i = 1; i < rightFieldBytesList.size(); i++) {
					byte[] rightFieldBytes = rightFieldBytesList.get(i);
					rightFieldBytes[3] = incrementPackageId();
					byteBuffer = sessionContext.writeToByteBuffer(rightFieldBytes, byteBuffer);
				}
				// write field eof
				EofPacket eofPacket = new EofPacket();
				eofPacket.packetId = incrementPackageId();
				byteBuffer = eofPacket.write(byteBuffer, sessionContext, true);
				sessionContext.write(byteBuffer);
				// LOGGER.info("header outputed ,packgId:" + eofPckg.packetId);
			} finally {
				writeLock.unlock();
			}
		}

	}

	public void writeHeader(List<byte[]> fieldBytesList) {
		if (headerWrited.compareAndSet(false, true)) {
			try {
				writeLock.lock();
				// write new header
				ResultSetHeaderPacket resultSetHeaderPacket = new ResultSetHeaderPacket();
				resultSetHeaderPacket.fieldCount = fieldBytesList.size();// -1;
				resultSetHeaderPacket.packetId = incrementPackageId();
				LOGGER.debug("packge id " + resultSetHeaderPacket.packetId);
				SessionContext sessionContext = session.getSessionContext();
				ByteBuffer byteBuffer = resultSetHeaderPacket.write(sessionContext.allocate(), sessionContext, true);
				// wirte a fields
				for (byte[] fieldBytes : fieldBytesList) {
					fieldBytes[3] = incrementPackageId();
					byteBuffer = sessionContext.writeToByteBuffer(fieldBytes, byteBuffer);
				}

				// write field eof
				EofPacket eofPacket = new EofPacket();
				eofPacket.packetId = incrementPackageId();
				byteBuffer = eofPacket.write(byteBuffer, sessionContext, true);
				sessionContext.write(byteBuffer);
				// LOGGER.info("header outputed ,packgId:" + eofPckg.packetId);
			} finally {
				writeLock.unlock();
			}
		}

	}

	public void writeRow(RowDataPacket rowDataPacket) {
		SessionContext sessionContext = session.getSessionContext();
		try {
			writeLock.lock();
			rowDataPacket.packetId = incrementPackageId();
			// 输出完整的 记录到客户端
			ByteBuffer byteBuffer = rowDataPacket.write(sessionContext.allocate(), sessionContext, true);
			sessionContext.write(byteBuffer);
			// LOGGER.info("write  row ,packgId:" + rowDataPkg.packetId);
		} finally {
			writeLock.unlock();
		}
	}

	public void writeEof() {
		SessionContext sessionContext = session.getSessionContext();
		EofPacket eofPacket = new EofPacket();
		eofPacket.packetId = incrementPackageId();
		ByteBuffer byteBuffer = eofPacket.write(sessionContext.allocate(), sessionContext, false);
		sessionContext.write(byteBuffer);
		LOGGER.info("write  eof ,packgId:" + eofPacket.packetId);
	}

	public Session getSession() {
		return session;
	}

	public void onTaskFinished(Task task) {

		boolean allFinished = tasks.taskFinished(task);
		if (allFinished && finished.compareAndSet(false, true)) {
			allTaskFinishedListener.onAllTaskFinished(this);
		}

	}
}
